import cv2
import os
import torch
from facenet_pytorch import MTCNN, InceptionResnetV1
import numpy as np
import pygame

# 获取设备
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)

# 初始化摄像头
camera_index = 0  # 默认使用内置摄像头
# 尝试使用外置摄像头
cap = cv2.VideoCapture(1)  # 外置摄像头通常为1
if not cap.isOpened():
    cap = cv2.VideoCapture(0)  # 如果外置摄像头不可用，则使用内置摄像头
if not cap.isOpened():
    raise IOError("无法打开摄像头，请检查摄像头连接或配置。")

# mtcnn模型加载【设置网络参数，进行人脸检测】
mtcnn = MTCNN(min_face_size=12, thresholds=[0.2, 0.2, 0.3], keep_all=False, device=device)

# InceptionResnetV1模型加载【用于获取人脸特征向量】
resnet = InceptionResnetV1(pretrained='vggface2').eval().to(device)

# 初始化pygame音频播放器
pygame.init()
pygame.mixer.init()

while True:
    ret, frame = cap.read()
    if not ret:
        print("无法获取帧数据，请检查摄像头连接或配置。")
        break

    # 使用mtcnn检测人脸
    face = mtcnn(frame)
    if face is not None:
        face = face.unsqueeze(0).to(device)
        with torch.no_grad():
            face_emb = resnet(face).detach().cpu()

        # 比较与tz文件夹中的所有npy文件
        known_embeddings_dir = 'tz'
        threshold = 0.8  # 阈值，根据实际情况调整
        is_known = False
        for filename in os.listdir(known_embeddings_dir):
            if filename.endswith('.npy'):
                known_emb = np.load(os.path.join(known_embeddings_dir, filename))
                known_emb_tensor = torch.from_numpy(known_emb).to(device)
                distance = (face_emb.to(device) - known_emb_tensor).norm().item()
                if distance < threshold:
                    is_known = True
                    print(f"与 {filename} 的人脸匹配，不是陌生人。")
                    break

        if not is_known:
            print("与所有已知人脸的距离都大于阈值，是陌生人。")
            # 播放报警音频（这里假设报警音频文件名为“alarm.wav”）
            pygame.mixer.music.load('alarm.wav')
            pygame.mixer.music.play(-1)  # -1表示循环播放
    else:
        print("未检测到人脸")

    # 按'q'键退出循环
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# 停止播放报警音频（如果正在播放）
pygame.mixer.music.stop()

# 释放资源
cap.release()
cv2.destroyAllWindows()
